#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
收费系统实现
支持免费模式、基础付费、专业付费、企业付费的完整计费系统
"""

import time
import json
import logging
from datetime import datetime, timedelta
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass, asdict
from enum import Enum
import sqlite3
import hashlib

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class UserTier(Enum):
    """用户等级枚举"""
    FREE = "free"
    BASIC = "basic"
    PRO = "pro"
    ENTERPRISE = "enterprise"

class TaskType(Enum):
    """任务类型枚举"""
    SINGLE_IMAGE = "single_image"
    BATCH_PROCESSING = "batch_processing"

@dataclass
class UserBilling:
    """用户计费信息"""
    user_id: str
    tier: UserTier
    monthly_quota: int
    used_quota: int
    daily_cost_limit: float
    current_daily_cost: float
    last_reset_date: datetime
    subscription_start: datetime
    subscription_end: Optional[datetime] = None
    is_active: bool = True

@dataclass
class TaskBilling:
    """任务计费信息"""
    task_id: str
    user_id: str
    task_type: TaskType
    image_count: int
    processing_time: float
    cost: float
    timestamp: datetime
    instance_cost: float = 0.0

class BillingSystem:
    """收费系统"""
    
    def __init__(self, db_path: str = "billing.db"):
        self.db_path = db_path
        self.pricing_config = self._load_pricing_config()
        self.cost_config = self._load_cost_config()
        self._init_database()
    
    def _load_pricing_config(self) -> Dict:
        """加载定价配置"""
        return {
            UserTier.FREE: {
                'monthly_fee': 0,
                'monthly_quota': 10,
                'overage_rate': 0.10,  # 超出费用 0.10元/张
                'batch_processing': False,
                'max_batch_size': 1,
                'queue_priority': 1,
                'daily_cost_limit': 0.5,
            },
            UserTier.BASIC: {
                'monthly_fee': 29,
                'monthly_quota': 500,
                'overage_rate': 0.08,
                'batch_processing': True,
                'max_batch_size': 10,
                'queue_priority': 2,
                'daily_cost_limit': 5.0,
            },
            UserTier.PRO: {
                'monthly_fee': 99,
                'monthly_quota': 2000,
                'overage_rate': 0.06,
                'batch_processing': True,
                'max_batch_size': 50,
                'queue_priority': 3,
                'daily_cost_limit': 15.0,
            },
            UserTier.ENTERPRISE: {
                'monthly_fee': 299,
                'monthly_quota': 10000,
                'overage_rate': 0.04,
                'batch_processing': True,
                'max_batch_size': 1000,
                'queue_priority': 4,
                'daily_cost_limit': 50.0,
            }
        }
    
    def _load_cost_config(self) -> Dict:
        """加载成本配置"""
        return {
            'hai_cost_per_second': 0.00033,  # 1.2元/小时
            'base_processing_cost': 0.01,    # 基础处理成本
            'batch_discount': 0.8,           # 批处理折扣
            'enterprise_discount': 0.7,      # 企业用户折扣
        }
    
    def _init_database(self):
        """初始化数据库"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        # 创建用户计费表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS user_billing (
                user_id TEXT PRIMARY KEY,
                tier TEXT NOT NULL,
                monthly_quota INTEGER NOT NULL,
                used_quota INTEGER DEFAULT 0,
                daily_cost_limit REAL NOT NULL,
                current_daily_cost REAL DEFAULT 0.0,
                last_reset_date TEXT NOT NULL,
                subscription_start TEXT NOT NULL,
                subscription_end TEXT,
                is_active BOOLEAN DEFAULT 1
            )
        ''')
        
        # 创建任务计费表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS task_billing (
                task_id TEXT PRIMARY KEY,
                user_id TEXT NOT NULL,
                task_type TEXT NOT NULL,
                image_count INTEGER NOT NULL,
                processing_time REAL NOT NULL,
                cost REAL NOT NULL,
                timestamp TEXT NOT NULL,
                instance_cost REAL DEFAULT 0.0,
                FOREIGN KEY (user_id) REFERENCES user_billing (user_id)
            )
        ''')
        
        # 创建成本监控表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS cost_monitoring (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                date TEXT NOT NULL,
                total_cost REAL NOT NULL,
                instance_cost REAL NOT NULL,
                processing_cost REAL NOT NULL,
                user_count INTEGER NOT NULL,
                task_count INTEGER NOT NULL
            )
        ''')
        
        conn.commit()
        conn.close()
    
    def create_user(self, user_id: str, tier: UserTier) -> UserBilling:
        """创建用户计费记录"""
        config = self.pricing_config[tier]
        
        billing = UserBilling(
            user_id=user_id,
            tier=tier,
            monthly_quota=config['monthly_quota'],
            used_quota=0,
            daily_cost_limit=config['daily_cost_limit'],
            current_daily_cost=0.0,
            last_reset_date=datetime.now().date(),
            subscription_start=datetime.now(),
            is_active=True
        )
        
        self._save_user_billing(billing)
        logger.info(f"Created user billing: {user_id} ({tier.value})")
        return billing
    
    def get_user_billing(self, user_id: str) -> Optional[UserBilling]:
        """获取用户计费信息"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('''
            SELECT * FROM user_billing WHERE user_id = ?
        ''', (user_id,))
        
        row = cursor.fetchone()
        conn.close()
        
        if row:
            return UserBilling(
                user_id=row[0],
                tier=UserTier(row[1]),
                monthly_quota=row[2],
                used_quota=row[3],
                daily_cost_limit=row[4],
                current_daily_cost=row[5],
                last_reset_date=datetime.fromisoformat(row[6]),
                subscription_start=datetime.fromisoformat(row[7]),
                subscription_end=datetime.fromisoformat(row[8]) if row[8] else None,
                is_active=bool(row[9])
            )
        return None
    
    def _save_user_billing(self, billing: UserBilling):
        """保存用户计费信息"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('''
            INSERT OR REPLACE INTO user_billing 
            (user_id, tier, monthly_quota, used_quota, daily_cost_limit, 
             current_daily_cost, last_reset_date, subscription_start, 
             subscription_end, is_active)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        ''', (
            billing.user_id, billing.tier.value, billing.monthly_quota,
            billing.used_quota, billing.daily_cost_limit, billing.current_daily_cost,
            billing.last_reset_date.isoformat(), billing.subscription_start.isoformat(),
            billing.subscription_end.isoformat() if billing.subscription_end else None,
            billing.is_active
        ))
        
        conn.commit()
        conn.close()
    
    def check_quota(self, user_id: str, task_type: TaskType, image_count: int = 1) -> Tuple[bool, str]:
        """检查用户配额"""
        billing = self.get_user_billing(user_id)
        if not billing:
            return False, "User not found"
        
        if not billing.is_active:
            return False, "User subscription inactive"
        
        # 检查月度配额
        if billing.used_quota + image_count > billing.monthly_quota:
            return False, f"Monthly quota exceeded ({billing.used_quota}/{billing.monthly_quota})"
        
        # 检查批处理限制
        config = self.pricing_config[billing.tier]
        if task_type == TaskType.BATCH_PROCESSING:
            if not config['batch_processing']:
                return False, "Batch processing not allowed for this tier"
            if image_count > config['max_batch_size']:
                return False, f"Batch size exceeds limit ({image_count}/{config['max_batch_size']})"
        
        # 检查每日成本限制
        estimated_cost = self._estimate_task_cost(billing.tier, task_type, image_count)
        if billing.current_daily_cost + estimated_cost > billing.daily_cost_limit:
            return False, f"Daily cost limit would be exceeded"
        
        return True, "Quota check passed"
    
    def _estimate_task_cost(self, tier: UserTier, task_type: TaskType, image_count: int) -> float:
        """估算任务成本"""
        config = self.pricing_config[tier]
        
        # 基础处理成本
        base_cost = self.cost_config['base_processing_cost'] * image_count
        
        # 批处理折扣
        if task_type == TaskType.BATCH_PROCESSING and image_count > 1:
            base_cost *= self.cost_config['batch_discount']
        
        # 企业用户折扣
        if tier == UserTier.ENTERPRISE:
            base_cost *= self.cost_config['enterprise_discount']
        
        # 实例成本估算 (基于处理时间)
        estimated_processing_time = image_count * 30  # 每张图片30秒
        instance_cost = estimated_processing_time * self.cost_config['hai_cost_per_second']
        
        return base_cost + instance_cost
    
    def calculate_task_cost(self, user_id: str, task_type: TaskType, image_count: int, 
                          processing_time: float, instance_time: float) -> float:
        """计算任务实际成本"""
        billing = self.get_user_billing(user_id)
        if not billing:
            return 0.0
        
        config = self.pricing_config[billing.tier]
        
        # 基础处理成本
        base_cost = self.cost_config['base_processing_cost'] * image_count
        
        # 批处理折扣
        if task_type == TaskType.BATCH_PROCESSING and image_count > 1:
            base_cost *= self.cost_config['batch_discount']
        
        # 企业用户折扣
        if billing.tier == UserTier.ENTERPRISE:
            base_cost *= self.cost_config['enterprise_discount']
        
        # 实例成本
        instance_cost = instance_time * self.cost_config['hai_cost_per_second']
        
        total_cost = base_cost + instance_cost
        
        # 如果超出月度配额，按超出费率计费
        if billing.used_quota >= billing.monthly_quota:
            overage_cost = image_count * config['overage_rate']
            total_cost += overage_cost
        
        return total_cost
    
    def record_task_billing(self, task_id: str, user_id: str, task_type: TaskType,
                          image_count: int, processing_time: float, instance_time: float) -> TaskBilling:
        """记录任务计费"""
        cost = self.calculate_task_cost(user_id, task_type, image_count, processing_time, instance_time)
        instance_cost = instance_time * self.cost_config['hai_cost_per_second']
        
        task_billing = TaskBilling(
            task_id=task_id,
            user_id=user_id,
            task_type=task_type,
            image_count=image_count,
            processing_time=processing_time,
            cost=cost,
            timestamp=datetime.now(),
            instance_cost=instance_cost
        )
        
        # 保存任务计费记录
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('''
            INSERT INTO task_billing 
            (task_id, user_id, task_type, image_count, processing_time, 
             cost, timestamp, instance_cost)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?)
        ''', (
            task_billing.task_id, task_billing.user_id, task_billing.task_type.value,
            task_billing.image_count, task_billing.processing_time, task_billing.cost,
            task_billing.timestamp.isoformat(), task_billing.instance_cost
        ))
        
        conn.commit()
        conn.close()
        
        # 更新用户计费信息
        self._update_user_usage(user_id, image_count, cost)
        
        logger.info(f"Recorded task billing: {task_id} - {user_id} - {cost:.4f}元")
        return task_billing
    
    def _update_user_usage(self, user_id: str, image_count: int, cost: float):
        """更新用户使用量"""
        billing = self.get_user_billing(user_id)
        if not billing:
            return
        
        # 检查是否需要重置每日成本
        today = datetime.now().date()
        if billing.last_reset_date < today:
            billing.current_daily_cost = 0.0
            billing.last_reset_date = today
        
        # 更新使用量
        billing.used_quota += image_count
        billing.current_daily_cost += cost
        
        self._save_user_billing(billing)
    
    def get_user_statistics(self, user_id: str) -> Dict:
        """获取用户统计信息"""
        billing = self.get_user_billing(user_id)
        if not billing:
            return {}
        
        # 获取任务统计
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('''
            SELECT 
                COUNT(*) as task_count,
                SUM(image_count) as total_images,
                SUM(cost) as total_cost,
                SUM(processing_time) as total_processing_time
            FROM task_billing 
            WHERE user_id = ? AND date(timestamp) >= date('now', '-30 days')
        ''', (user_id,))
        
        stats = cursor.fetchone()
        conn.close()
        
        config = self.pricing_config[billing.tier]
        
        return {
            'user_id': user_id,
            'tier': billing.tier.value,
            'monthly_quota': billing.monthly_quota,
            'used_quota': billing.used_quota,
            'remaining_quota': billing.monthly_quota - billing.used_quota,
            'quota_usage_percent': (billing.used_quota / billing.monthly_quota) * 100,
            'daily_cost_limit': billing.daily_cost_limit,
            'current_daily_cost': billing.current_daily_cost,
            'daily_cost_usage_percent': (billing.current_daily_cost / billing.daily_cost_limit) * 100,
            'last_30_days': {
                'task_count': stats[0] or 0,
                'total_images': stats[1] or 0,
                'total_cost': stats[2] or 0.0,
                'total_processing_time': stats[3] or 0.0,
            },
            'subscription_status': 'active' if billing.is_active else 'inactive',
            'subscription_start': billing.subscription_start.isoformat(),
            'subscription_end': billing.subscription_end.isoformat() if billing.subscription_end else None,
        }
    
    def get_system_revenue_report(self, start_date: datetime, end_date: datetime) -> Dict:
        """获取系统收入报告"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        # 获取收入统计
        cursor.execute('''
            SELECT 
                ub.tier,
                COUNT(DISTINCT ub.user_id) as user_count,
                SUM(tb.cost) as total_revenue,
                SUM(tb.image_count) as total_images,
                COUNT(tb.task_id) as total_tasks
            FROM user_billing ub
            LEFT JOIN task_billing tb ON ub.user_id = tb.user_id 
                AND date(tb.timestamp) BETWEEN date(?) AND date(?)
            WHERE ub.is_active = 1
            GROUP BY ub.tier
        ''', (start_date.isoformat(), end_date.isoformat()))
        
        tier_stats = cursor.fetchall()
        
        # 获取成本统计
        cursor.execute('''
            SELECT 
                SUM(instance_cost) as total_instance_cost,
                SUM(cost - instance_cost) as total_processing_cost,
                SUM(cost) as total_cost
            FROM task_billing 
            WHERE date(timestamp) BETWEEN date(?) AND date(?)
        ''', (start_date.isoformat(), end_date.isoformat()))
        
        cost_stats = cursor.fetchone()
        conn.close()
        
        # 计算订阅收入
        subscription_revenue = 0
        for tier, user_count, _, _, _ in tier_stats:
            tier_enum = UserTier(tier)
            monthly_fee = self.pricing_config[tier_enum]['monthly_fee']
            subscription_revenue += user_count * monthly_fee
        
        total_revenue = sum(stat[2] or 0 for stat in tier_stats)
        total_cost = cost_stats[2] or 0 if cost_stats else 0
        net_profit = total_revenue - total_cost
        
        return {
            'period': {
                'start_date': start_date.isoformat(),
                'end_date': end_date.isoformat(),
            },
            'revenue': {
                'subscription_revenue': subscription_revenue,
                'usage_revenue': total_revenue,
                'total_revenue': subscription_revenue + total_revenue,
            },
            'costs': {
                'instance_cost': cost_stats[0] or 0 if cost_stats else 0,
                'processing_cost': cost_stats[1] or 0 if cost_stats else 0,
                'total_cost': total_cost,
            },
            'profit': {
                'gross_profit': total_revenue,
                'net_profit': net_profit,
                'profit_margin': (net_profit / total_revenue * 100) if total_revenue > 0 else 0,
            },
            'usage': {
                'total_tasks': sum(stat[4] or 0 for stat in tier_stats),
                'total_images': sum(stat[3] or 0 for stat in tier_stats),
            },
            'tier_breakdown': [
                {
                    'tier': stat[0],
                    'user_count': stat[1],
                    'revenue': stat[2] or 0,
                    'images': stat[3] or 0,
                    'tasks': stat[4] or 0,
                }
                for stat in tier_stats
            ]
        }
    
    def check_daily_cost_limits(self) -> List[Dict]:
        """检查每日成本限制"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('''
            SELECT user_id, tier, current_daily_cost, daily_cost_limit
            FROM user_billing 
            WHERE is_active = 1 AND current_daily_cost > daily_cost_limit * 0.8
        ''')
        
        alerts = []
        for row in cursor.fetchall():
            user_id, tier, current_cost, limit = row
            usage_percent = (current_cost / limit) * 100
            
            alert_level = 'warning'
            if usage_percent >= 100:
                alert_level = 'critical'
            elif usage_percent >= 95:
                alert_level = 'emergency'
            
            alerts.append({
                'user_id': user_id,
                'tier': tier,
                'current_cost': current_cost,
                'limit': limit,
                'usage_percent': usage_percent,
                'alert_level': alert_level,
            })
        
        conn.close()
        return alerts

# 使用示例
def main():
    """使用示例"""
    billing = BillingSystem()
    
    # 创建用户
    user_id = "test_user_001"
    billing.create_user(user_id, UserTier.PRO)
    
    # 检查配额
    can_process, reason = billing.check_quota(user_id, TaskType.SINGLE_IMAGE, 1)
    print(f"Can process: {can_process}, Reason: {reason}")
    
    # 记录任务计费
    task_billing = billing.record_task_billing(
        task_id="task_001",
        user_id=user_id,
        task_type=TaskType.SINGLE_IMAGE,
        image_count=1,
        processing_time=30.0,
        instance_time=45.0
    )
    print(f"Task billing: {task_billing}")
    
    # 获取用户统计
    stats = billing.get_user_statistics(user_id)
    print(f"User stats: {stats}")
    
    # 获取收入报告
    start_date = datetime.now() - timedelta(days=30)
    end_date = datetime.now()
    report = billing.get_system_revenue_report(start_date, end_date)
    print(f"Revenue report: {report}")

if __name__ == '__main__':
    main()
